import re,csv,os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import confusion_matrix
import seaborn as sns
import util

def extractData(path:str):
    with open(path) as fp:
        data=fp.read()
    pattern=r"accuracy:\s*([-+]?\d*.\d+|\d+)\nprecision:\s*([-+]?\d*.\d+|\d+)\nrecall:\s*([-+]?\d*.\d+|\d+)\nf1:\s*([-+]?\d*.\d+|\d+)"
    found= re.findall(pattern,data)
    return [found[i] for i in range(len(found)) if i%2==0],[found[i] for i in range(len(found)) if i%2!=0]

def saveCsv(data,path:str):
    with open(path, "w", newline="", encoding="utf-8") as file:
        writer = csv.writer(file)
        writer.writerows(data)

def plot_pr_comparison(file_paths: list, 
                       model_names: tuple = ("Train whole model", "Train fc only"),
                       output_path: str = "pr_comparison.png",
                       figsize: tuple = (16, 6),
                       dpi: int = 300) -> None:
    """
    专用于精度(Precision)和召回率(Recall)的可视化对比
    - file_paths: 顺序为 [模型1训练数据, 模型1验证数据, 模型2训练数据, 模型2验证数据]
    """
    # 1. 数据验证与读取
    for path in file_paths:
        if not os.path.exists(path):
            raise FileNotFoundError(f"文件路径不存在: {os.path.abspath(path)}")
    
    # 2. 数据整合
    metric_columns = ['Precision', 'Recall']  # 仅保留两个指标
    datasets = ["on train dataset", "on validate dataset", "on train dataset", "on validate dataset"]
    
    dfs = []
    for i, path in enumerate(file_paths):
        model = model_names[0] if i < 2 else model_names[1]
        df = pd.read_csv(path, header=None, names=[metric_columns[0], metric_columns[1]])  # 仅读取第2、3列
        df['Epoch'] = range(len(df))
        df['Model'] = model
        df['Dataset'] = datasets[i]
        dfs.append(df[['Epoch', 'Model', 'Dataset', 'Precision', 'Recall']])  # 过滤无关列
    
    combined_df = pd.concat(dfs)

    # 3. 可视化配置
    plt.style.use('seaborn-v0_8-whitegrid')
    plt.rcParams.update({'font.size': 12})
    
    # 4. 创建画布（调整为横向布局）
    fig, axs = plt.subplots(1, 2, figsize=figsize, dpi=dpi)
    colors = {model_names[0]: '#1f77b4', model_names[1]: '#ff7f0e'}
    line_styles = {'on train dataset': '-', 'on validate dataset': '--'}

    # 5. 绘制子图
    for idx, metric in enumerate(metric_columns):
        ax = axs[idx]
        for (model, dataset), group in combined_df.groupby(['Model', 'Dataset']):
            ax.plot(group['Epoch'], 
                    group[metric],
                    color=colors[model],
                    linestyle=line_styles[dataset],
                    linewidth=1,
                    alpha=0.8,
                    label=f"{model} {dataset}")
            
            # ax.set_title(f'{metric} Comparison', fontsize=14)
            ax.set_ylim(0.2,1.01)
            ax.set_xlabel('Epoch', fontsize=10)
            ax.set_ylabel(metric, fontsize=10)
            ax.grid(True, alpha=0.3)
            ax.legend(loc='lower right')

    # 6. 全局配置
    plt.tight_layout()
    plt.savefig(output_path, bbox_inches='tight')
    plt.close()

def extract():
    validate,train=extractData("log/train_full/train_full.log")
    saveCsv(validate,"validate_data.csv")
    saveCsv(train,"tran_data.csv")
    print(validate,train)


def plot_confusion_matrix(y_true, y_pred, 
                          class_names=None, 
                          normalize=False,
                          title="Confusion Matrix",
                          cmap="Blues",
                          figsize=(16, 12),
                          save_path=None,
                          **kwargs):
    """
    用 Seaborn 绘制混淆矩阵

    参数:
    - y_true: 真实标签 (array-like)
    - y_pred: 预测标签 (array-like)
    - class_names: 类别名称列表 (list)，默认自动推断
    - normalize: 是否显示百分比 (bool)
    - title: 图表标题 (str)
    - cmap: 颜色映射 (str)
    - figsize: 图表尺寸 (tuple)
    - save_path: 保存路径 (str, 可选)
    - **kwargs: 其他 seaborn.heatmap 参数
    """
    # 计算混淆矩阵
    cm = confusion_matrix(y_true, y_pred)
    
    # 自动生成类别名称（如果没有提供）
    if class_names is None:
        class_names = [f"Class {i}" for i in np.unique(y_true)]
    
    # 标准化混淆矩阵
    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
        fmt = '.2%'
    else:
        fmt = 'd'

    # 创建画布
    plt.figure(figsize=figsize)
    
    # 绘制热力图
    ax = sns.heatmap(
        cm,
        annot=True,
        fmt=fmt,
        cmap=cmap,
        xticklabels=class_names,
        yticklabels=class_names,
        **kwargs
    )

    # 美化样式
    # ax.set_title(title, fontsize=14, pad=20)
    ax.set_xlabel("Predicted Label", fontsize=12)
    ax.set_ylabel("True Label", fontsize=12)
    
    # 调整刻度标签（多分类时旋转45度）
    if len(class_names) > 5:
        plt.xticks(rotation=45, ha='right')
        plt.yticks(rotation=0)
    else:
        plt.xticks(rotation=0)
        plt.yticks(rotation=0)

    # 保存图片
    if save_path:
        plt.savefig(save_path, bbox_inches='tight', dpi=300)
        print(f"Confusion matrix saved to: {save_path}")

    # plt.show()



if __name__=="__main__":
    # plot_pr_comparison(["log/train_full/train_full.csv","log/train_full/validate_full.csv","log/train_fc/train_fc.csv","log/train_fc/validate_fc.csv"])
    prediction=np.load("result/svm_poly.npy")
    _,true_labels=util.load_np_data("test.npz")
    class_names=["Bean",  "Bitter_Gourd",  "Bottle_Gourd",  "Brinjal",  "Broccoli",  "Cabbage",  "Capsicum",  "Carrot",  "Cauliflower",  "Cucumber",  "Papaya",  "Potato",  "Pumpkin",  "Radish",  "Tomato"]
    plot_confusion_matrix(true_labels,prediction,class_names,save_path="confusion.png")